import pandas as pd
import networkx as nx
import re

# '62<TCPv6:[[::ffff:10.0.2.16]:51352->[::ffff:46.82.174.69]:443]>,{sa_family=AF_INET6,sin6_port=htons(51352),inet_pton(AF_INET6,"::ffff:10.0.2.16",&sin6_addr),sin6_flowinfo=htonl(0),sin6_scope_id=0},[128->28]'
def find_ip(params:str) -> tuple[list,bool]:
    results = []
    """  
    match connect() parameters
        ex: ...htons(443),....::ffff:40.70.161.102...
        return port and ip
        ipv4 or ipv6
    """
    connect_ipv4_pattern = r'htons\((\d+)\).*?::ffff:(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})'
    connect_ipv6_pattern = r'htons\((\d+)\).*?inet_pton\(AF_INET6,"([\w:]+)"'

    """
    match connected socket
        ex:63<TCPv6:[[::ffff:10.0.2.16]:39830->[::ffff:104.18.215.59]:443]>
        return source ip\port  des ip\port
    """
    compl_ipv4_pattern = r'::ffff:(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\]:(\d+).*?::ffff:(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\]:(\d+)'
    """ 
    match ipv6 
        ex:<UDPv6:[[fec0::41dd:35fd:9649:f8cf]:51994->[2001:4860:4860::8888]:53]>
    """
    compl_ipv6_pattern = r'\[(\w{0,4}(:\w{0,4}){1,7})\]:(\d+)->\[(\w{0,4}(:\w{0,4}){1,7})\]:(\d+)'

    """  
    match content of sendmsg \ revfrom 
        but it is uncompleted
        return portion of a load
    """
    match_connect_ipv4_pattern = re.search(connect_ipv4_pattern,params)
    if match_connect_ipv4_pattern:
        port = match_connect_ipv4_pattern.group(1)
        ip = match_connect_ipv4_pattern.group(2)
        results.append((port,ip))
        # return results,True
    match_connect_ipv6_pattern = re.search(connect_ipv6_pattern,params)
    if match_connect_ipv6_pattern:
        port = match_connect_ipv6_pattern.group(1)
        ip = match_connect_ipv6_pattern.group(2)
        results.append((port,ip))
        # return results,True

    match_compl_ipv4 = re.search(compl_ipv4_pattern, params)
    if match_compl_ipv4:
        src_ip = match_compl_ipv4.group(1)
        src_port = match_compl_ipv4.group(2)
        dest_ip = match_compl_ipv4.group(3)
        dest_port = match_compl_ipv4.group(4)
        results.append((src_port, src_ip))
        results.append((dest_port, dest_ip))

        # return results,True
    match_compl_ipv6 = re.search(compl_ipv6_pattern, params)
    if match_compl_ipv6:
        src_ip = match_compl_ipv6.group(1)
        src_port = match_compl_ipv6.group(3)
        dest_ip = match_compl_ipv6.group(4)
        dest_port = match_compl_ipv6.group(6)
        results.append((src_port, src_ip))
        results.append((dest_port, dest_ip))
        # return results,True
    
    if len(results)>0:
        return results,True
    return '',False

def build_fd_graph(inp, output_gexf_path):
    """
    Generate a function call graph based on file descriptors and save it as GEXF.
    """
    df = pd.read_csv(inp)
    fd_pattern = re.compile(r"(\d+<.*?>)")
    param_fds_pattern = re.compile(r"(\d+<.*?>)")
    connect_pattern = re.compile(r'\d+<(?:TCP|UDP|TCPv6|UDPv6):\[(?:\[(?:[0-9a-fA-F:\.]+)\]:(?:\d+)(?:->)?(?:\[(?:[0-9a-fA-F:\.]+)\]:(?:\d+))?|(?:[0-9\.]+):(?:\d+)(?:->(?:[0-9\.]+):(?:\d+))?)\]>'\
                                 r'|'\
                                 r'\d+<.*?>'
                                 )
    G = nx.DiGraph()
    fd_dict = {}
    connect_info = {}  # Cache to store the connect information with file descriptors
    traffic = 1
    for i, row in df.iterrows():
        if '-1' in str(row['Result']) and str(row['Function']) != 'connect':
            continue

        result_fd = fd_pattern.search(str(row['Result']))
        param_fds = param_fds_pattern.findall(str(row['Params']))
        connect_fds =param_fds_pattern.findall(str(row['Params']))
        ip_info, valid = find_ip(str(row['Params']))
        now_fd = ''
        if result_fd:
            """  
                结果有fd，则直接添加进fd_dict，并且直接添加到图里。
            """
            full_fd = result_fd.group(0)
            params_str = str(row['Params'])
            if full_fd not in G:
                G.add_node(full_fd, params=params_str)
            else:
                # 如果节点已存在，将新参数添加到params列表中
                if 'params' not in G.nodes[full_fd]:
                    G.nodes[full_fd]['params'] = params_str
                else:
                    G.nodes[full_fd]['params'] += f" {params_str}"
            fd_dict[full_fd] = full_fd
        
        """
            做好connect的映射
        """
        if valid:
            # 如果为connect则建立映射。
            if str(row['Function']) == 'connect':
                if param_fds:
                    full_fd = param_fds[0]
                    connect_info[full_fd] = ip_info
            else:
                for key,value_list in connect_info.items():
                    for ip in ip_info:
                        for value in value_list :
                            if ip == value:
                                now_fd = key
        # return now_fd ：已经经过转化的fd
        
        first_match = None
        for fd in param_fds:
            # 如果第一个是因为ip信息无法识别的fd.
            if (first_match is None) and ('TCP' in fd or 'UDP' in fd) and (fd not in fd_dict):
                if now_fd != '':
                    current_node = fd_dict[now_fd]
                    new_node = f"{now_fd}-{row['Function']}"
                    params_str = str(row['Params'])
                    if new_node not in G:
                        G.add_node(new_node, params=params_str)
                        G.add_edge(current_node,new_node)
                    else:
                        G.nodes[new_node]['params'] += f" {params_str}"
                        G.add_edge(current_node,new_node)
                    fd_dict[now_fd] = new_node
                    first_match = now_fd
                    continue
            
            if first_match != None and now_fd !='' and ('TCP' in fd or 'UDP' in fd):
                G.add_edge(now_fd,fd_dict[first_match])
                continue
            
            if fd not in fd_dict:
                continue
            
            # 如果是第一次匹配，而且在fd_dict，就直接弄。
            if first_match is None:
                first_match = fd
                current_node = fd_dict[fd]
                new_node = f"{fd}-{row['Function']}"
                params_str = str(row['Params'])
                if new_node not in G:
                    G.add_node(new_node, params=params_str)
                else:
                    G.nodes[new_node]['params'] += f" {params_str}"
                G.add_edge(current_node, new_node)
                fd_dict[fd] = new_node
            # 如果不是第一次了
            else:
                if now_fd != '':
                    G.add_edge(now_fd,fd_dict[first_match])
                else:
                    G.add_edge(fd, fd_dict[first_match])
        
        # 如果Matched_Brust 不为空:
        # 添加一个节点

        # 这里还没完成映射.
        if result_fd and param_fds:
            for param in param_fds:
                if ('TCP' in fd or 'UDP' in fd) and now_fd != '':
                    if now_fd in fd_dict and result_fd.group(0) in fd_dict:
                        G.add_edge(now_fd, result_fd.group(0))
                else:
                    if now_fd in fd_dict and result_fd.group(0) in fd_dict:
                        G.add_edge(param, result_fd.group(0))
        
        if pd.notna(row['Matched_Brust']):
            if str(row['Function']) in ['sendmsg', 'sendto']:
                sn = f"{now_fd}-{row['Function']}"
                tn = f'traffic{str(traffic)}'
                traffic = traffic + 1
                G.add_node(traffic,burst = str(row['Matched_Brust']))
                G.add_edge(sn,tn)
            if str(row['Function']) in ['recvfrom', 'recvmsg']:    
                sn = f"{now_fd}-{row['Function']}"
                tn = f'traffic{str(traffic)}'
                traffic = traffic + 1
                G.add_node(traffic,burst = str(row['Matched_Brust']))
                G.add_edge(tn,sn)
    if G.number_of_nodes() > 0:
        nx.write_gexf(G, output_gexf_path)



build_fd_graph('/dataset/goodware/a8dfb53f36c703e1faecfd0f189e7cc8/a8dfb53f36c703e1faecfd0f189e7cc8.csv','test.gexf')